"""
修复UUID问题的迁移脚本
"""

import os
import sys
import asyncio
import json
import uuid
from pathlib import Path

# 设置环境变量
os.environ['SUPABASE_URL'] = 'https://wfmarihxfyyqljzaakbk.supabase.co'
os.environ['SUPABASE_ANON_KEY'] = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6IndmbWFyaWh4Znl5cWxqemFha2JrIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NTcxNTM5MDQsImV4cCI6MjA3MjcyOTkwNH0.9KoJUGhMbMNyNZhFHaqUIv3mt4z26XnbN4sMq0XlUxo'

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

DEFAULT_USER_ID = '00000000-0000-0000-0000-000000000000'

def generate_uuid_from_string(s):
    """从字符串生成固定的UUID"""
    return str(uuid.uuid5(uuid.NAMESPACE_DNS, s))

class FixedMigrator:
    def __init__(self, data_dir):
        self.data_dir = Path(data_dir)
        self.stats = {
            'sessions_migrated': 0,
            'messages_migrated': 0,
            'roles_migrated': 0,
            'errors': []
        }
        self.db_service = None
        self.id_mapping = {}  # 原ID到新UUID的映射
        
    async def migrate(self):
        print("=== 修复UUID问题的迁移 ===")
        
        try:
            from database import get_db_service
            self.db_service = get_db_service()
            
            if not self.db_service.initialize():
                print("[ERROR] 数据库初始化失败")
                return False
                
            print("[OK] 数据库连接成功")
            
            # 先迁移角色（因为会话可能依赖角色）
            await self._migrate_roles()
            
            # 再迁移会话和消息
            await self._migrate_sessions()
            
            self._show_results()
            
            return True
            
        except Exception as e:
            print(f"[ERROR] 迁移失败: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    async def _migrate_roles(self):
        print("\\n迁移角色数据...")
        
        roles_file = self.data_dir / "roles.json"
        if not roles_file.exists():
            print("[WARN] 角色文件不存在")
            return
            
        try:
            with open(roles_file, 'r', encoding='utf-8') as f:
                roles_data = json.load(f)
                
            print(f"发现 {len(roles_data)} 个角色")
            
            for role_id, role_data in roles_data.items():
                try:
                    # 生成新的UUID
                    new_uuid = generate_uuid_from_string(f"role_{role_id}")
                    self.id_mapping[role_id] = new_uuid
                    
                    # 转换角色特征
                    from database import RoleTrait
                    traits = []
                    if 'custom_traits' in role_data:
                        for trait in role_data.get('custom_traits', []):
                            if isinstance(trait, dict) and 'key' in trait and 'value' in trait:
                                traits.append(RoleTrait(key=trait['key'], value=trait['value']))
                    
                    # 创建角色（不传入id，让数据库自动生成）
                    role = await self.db_service.create_custom_role(
                        user_id=DEFAULT_USER_ID,
                        name=role_data.get('name', ''),
                        title=role_data.get('title', ''),
                        input_spec=role_data.get('input_spec', ''),
                        output_spec=role_data.get('output_spec', ''),
                        custom_traits=traits,
                        system_prompt_snapshot=role_data.get('system_prompt_snapshot', ''),
                        usage_count=role_data.get('usage_count', 0)
                    )
                    
                    if role:
                        print(f"[OK] 角色: {role.name}")
                        self.stats['roles_migrated'] += 1
                        # 更新映射为实际创建的UUID
                        self.id_mapping[role_id] = role.id
                    else:
                        print(f"[ERROR] 角色创建失败: {role_data.get('name')}")
                        
                except Exception as e:
                    print(f"[ERROR] 角色迁移出错: {e}")
                    self.stats['errors'].append(f"角色错误: {e}")
                    
        except Exception as e:
            print(f"[ERROR] 读取角色文件失败: {e}")
            
    async def _migrate_sessions(self):
        print("\\n迁移会话数据...")
        
        sessions_file = self.data_dir / "sessions.json"
        if not sessions_file.exists():
            print("[WARN] 会话文件不存在")
            return
            
        try:
            with open(sessions_file, 'r', encoding='utf-8') as f:
                sessions_data = json.load(f)
                
            print(f"发现 {len(sessions_data)} 个会话")
            
            for session_id, session_data in sessions_data.items():
                try:
                    # 生成新的UUID
                    new_uuid = generate_uuid_from_string(f"session_{session_id}")
                    
                    # 处理角色ID映射
                    role_id = session_data.get('roleId')
                    if role_id and role_id in self.id_mapping:
                        role_id = self.id_mapping[role_id]
                    
                    # 创建会话（不传入id，让数据库自动生成）
                    session = await self.db_service.create_session(
                        user_id=DEFAULT_USER_ID,
                        title=session_data.get('title', '新会话'),
                        role_id=role_id,
                        model=session_data.get('model', 'deepseek-v3-1-250821'),
                        search_enabled=session_data.get('searchEnabled', False),
                        memo=session_data.get('memo'),
                        is_favorite=session_data.get('isFavorite', False),
                        is_archived=session_data.get('isArchived', False)
                    )
                    
                    if not session:
                        print(f"[ERROR] 会话创建失败: {session_data.get('title')}")
                        continue
                        
                    print(f"[OK] 会话: {session.title}")
                    self.stats['sessions_migrated'] += 1
                    
                    # 迁移消息
                    messages = session_data.get('messages', [])
                    for i, message_data in enumerate(messages):
                        try:
                            message = await self.db_service.add_message(
                                session_id=session.id,
                                role=message_data.get('role', 'user'),
                                content=message_data.get('content', ''),
                                model=message_data.get('model'),
                                message_index=i + 1
                            )
                            
                            if message:
                                self.stats['messages_migrated'] += 1
                                
                        except Exception as e:
                            print(f"[ERROR] 消息迁移失败: {e}")
                            self.stats['errors'].append(f"消息错误: {e}")
                            
                except Exception as e:
                    print(f"[ERROR] 会话迁移失败: {e}")
                    self.stats['errors'].append(f"会话错误: {e}")
                    
        except Exception as e:
            print(f"[ERROR] 读取会话文件失败: {e}")
    
    def _show_results(self):
        print("\\n=== 迁移完成 ===")
        print(f"会话: {self.stats['sessions_migrated']} 个")
        print(f"消息: {self.stats['messages_migrated']} 条")
        print(f"角色: {self.stats['roles_migrated']} 个")
        
        if self.stats['errors']:
            print(f"\\n错误: {len(self.stats['errors'])} 个")
            for error in self.stats['errors'][:3]:
                print(f"  - {error}")
        else:
            print("\\n[SUCCESS] 迁移成功，没有错误！")

async def main():
    migrator = FixedMigrator("./data")
    await migrator.migrate()

if __name__ == "__main__":
    asyncio.run(main())